
!-------------------------------------------------------------------------------------
!
!    call evolve_F_long_advect( dfu, dt)
!
!    call evolve_F_trans_advect( dfu, Fbgk, uxMat, dt)
!
!    call evolve_efield( dfu, Ex2D, Ez2D)
!
!    call evolve_F_accel( dfu, Ebgk, DVbgk, Ez2D, dt)
!
!
!   call evolve_U_long_advect( uxMat, dt)
!
!   call evolve_U_trans_advect( uxMat, dt)
!
!   call evolve_U_sourceE( uxMat, Ex2D, dt)



!-------------------------------------------------------------------------------------
subroutine evolve_lin_F_long_advect(dfu, dt)
!-------------------------------------------------------------------------------------

  use dimensions
  use parameters, only: iunit, Dzzzz, SideLoss
  use fft_wrap

  implicit none

  real,    intent(inout) :: dfu(jv, jz, jy, jx, jDir)
  real,    intent(in)    :: dt

  integer                :: iv, ix, iy, iDir
  complex                :: fz(jz)
  complex                :: iVdt


  !-- advect in z-direction

  iy=jy

  do iDir=1,jDir
     do ix=1, jx
        do iv=1,jv
           
           iVdt = iunit * dt * vVec(iv)

           fz(:) = dfu(iv, :, iy, ix, iDir)
           call fft_forward(jz,fz,fz)
           fz = fz/float(jz)
           fz(:) = fz(:) * exp ( -iVdt * kZvec(:) - Dzzzz * kZvec(:)**4 *dt - SideLoss*dt )
           call fft_backward(jz,fz,fz)
           dfu(iv, :, iy, ix, iDir) = fz(:)

        enddo
     enddo
  enddo

end subroutine evolve_lin_F_long_advect


!-------------------------------------------------------------------------------------
subroutine evolve_lin_F_trans_advect(dfu, Fbgk, uxMat, dt)
!-------------------------------------------------------------------------------------

  use dimensions
  use parameters, only: iunit, ux0
  use fft_wrap

  implicit none

  real,    intent(inout) :: dfu(jv, jz, jy, jx, jDir)
  real,    intent(in)    :: Fbgk(jv, jz)
  real,    intent(in)    :: uxMat(jz, jy, jx, jDir)
  real,    intent(in)    :: dt


  complex                :: fx(jx), ux(jx), b(jx), c(jx), e(jx)

  integer                :: ix, iz, iy, iv, iDir


  !-- advect in x-direction

  iy=jy

  do iDir=1,jDir
 
      b = -iunit * ux0(iDir) *  kXvec(:)
      e =  exp(b*dt) 
 
      do iz=1, jz

        ux(:) = uxMat(iz,iy,:,iDir)
        call fft_forward(jx,ux,ux)
        ux = ux/float(jx)

        do iv=1,jv
           
           c = Fbgk(iv,iz) * ux(:) / ux0(iDir)

           fx(:) = dfu(iv, iz, iy, :, iDir)
           call fft_forward(jx,fx,fx)
           fx = fx/float(jx)

           fx  = fx*e + c * (e - 1)

           call fft_backward(jx,fx,fx)
           dfu(iv, iz, iy, :, iDir) = fx(:)
  
        enddo
     enddo
  enddo

end subroutine evolve_lin_F_trans_advect


!-------------------------------------------------------------------------------------
subroutine evolve_lin_F_accel(dfu, Ebgk, DVbgk, Ez2D, dt)
!-------------------------------------------------------------------------------------

  use dimensions
  use parameters, only: iunit, Dvvvv
  use fft_wrap

  implicit none

  real,    intent(inout) :: dfu(jv, jz, jy, jx, jDir)
  real,    intent(in)    :: Ebgk(jz)
  real,    intent(in)    :: DVbgk(jv, jz)
  real,    intent(in)    :: Ez2D(jz, jy, jx)
  real,    intent(in)    :: dt

  integer                :: iz, iv, ix, iy, iDir
  complex                :: fv(jv), fb(jv), fa(jv) 
  complex                :: b(jv), c(jv), e(jv)


  !-- advect in v-direction

  iy = jy

  do idir=1,jDir
     do ix=1, jx

        do iv=1,jv
           fa(iv) = sum(dfu(iv,:,iy,ix,iDir))/jz;
        enddo
   
        call fft_forward(jv,fa,fa)
        fa=fa/float(jv)

        do iz=1,jz

           fb(:) = DVbgk(:,iz);
           call fft_forward(jv,fb,fb)
           fb = fb/float(jv)

           b(:) = - iunit*Ebgk(iz)*kVvec(:) - Dvvvv*kVvec(:)**4
           e(:) = exp(b*dt)

           c(:) = - fb(:)*Ez2D(iz,iy,ix)  + Dvvvv*kVvec(:)**4*fa
           c(:) = c/b
           c(1) = 0

           fv(:) = dfu(:, iz, iy, ix, iDir)
           call fft_forward(jv,fv,fv)
           fv = fv/float(jv)

           fv(:)  = fv*e + c*(e-1)
    
           call fft_backward(jv,fv,fv)
           dfu(:, iz, iy, ix, iDir) = fv(:)

        end do
     end do
  end do
 
end subroutine evolve_lin_F_accel


!-------------------------------------------------------------------------------------
subroutine evolve_lin_F_combined(dfu, DVbgk, ftEz, dt)
!-------------------------------------------------------------------------------------

  use dimensions
  use parameters, only: iunit, Dzzzz
  use fft_wrap

  implicit none

  real,    intent(inout) :: dfu(jv, jz, jy, jx, jDir)
  real,    intent(in)    :: DVbgk(jv, jz)
  complex, intent(in)    :: ftEz(jz, jy, jx)
  real,    intent(in)    :: dt

  integer                :: iv, ix, iy, iz, iDir
  complex                :: fz(jz), fzAdd(jz)
  complex                :: expV(jz), expDzzzz(jz) 
  complex                :: iKzV(jz)


  !-- advect in z-direction


  iy=jy

  expDzzzz(:) =  exp(-Dzzzz * kZvec(:)**4 *dt)

  do iDir=1,jDir
     do ix=1, jx
        do iv=1,jv
           
           fz(:) = dfu(iv, :, iy, ix, iDir)

           call fft_forward(jz,fz,fz)
           fz = fz/float(jz)

           iKzV(:) = iunit * kZvec(:) * vVec(iv)

           expV(:) = exp ( -iKzV(:)*dt )


           fzAdd(:) =  - ftEz(:,iy,ix)*DVbgk(iv,:) * (expV(:) - 1) / (-iKzV(:))
           fzAdd(1) =  - ftEz(1,iy,ix)*DVbgk(iv,1) * dt

           fz(:) = fz(:)*expV(:)*expDzzzz(:) + fzAdd(:)
           !fz(:) = fz(:) + fzAdd(:)

           call fft_backward(jz,fz,fz)
           dfu(iv, :, iy, ix, iDir) = fz(:)

        enddo
     enddo
  enddo


end subroutine evolve_lin_F_combined


!*************************************************************************************

!-------------------------------------------------------------------------------------
subroutine evolve_lin_efield(dfu, Ex3D, Ez3D, ftEz, filter)
!-------------------------------------------------------------------------------------

  use dimensions
  use parameters, only: iunit, kz, phi0
  use fft_wrap
  use fft_3d
  use mpi_wrap

  implicit none

  real,    intent(in)   :: dfu(jv, jz, jy, jx, jDir)
  real,    intent(out)  :: Ex3D(jz, jy, jx)
  real,    intent(out)  :: Ez3D(jz, jy, jx)
  complex, intent(out)  :: ftEz(jz, jy, jx)
  real,    intent(in)   :: filter(jz,jy,jx)

  real                  :: tmp(jz, jy, jx)
  real                  :: dens(jz, jy, jx)
  real,    save         :: kk3D(jz, jy, jx)

  complex               :: dummy3D(jz, jy, jx)
  complex               :: phiFT(jz, jy, jx)
  complex, save         :: phiExtFT(jz)

  integer               :: ix, iz, iy, iv, idir
  logical, save         :: first_call = .true.

! real, save            :: weights(jv)


  !-- generate matrix of k-squared

  if (first_call) then

     phiExtFT(:) = phi0 *exp ( iunit * kz * zvec(:) ) 
     call fft_forward(jz,phiExtFT,phiExtFT)
     phiExtFT = phiExtFT/float(jz)
    
     do ix=1,jx
        do iy=1,jy
           do iz=1,jz
              kk3D(iz,iy,ix) = kXvec(ix)**2 + kZvec(iz)**2  + kYvec(iy)**2
           end do
        end do
     end do

     !do iv=1,jv
        !weights(iv) = ( mod(iv,2) + 1 ) * dv * 2/3      ! simpson
        !weights(iv) = ( abs(mod(iv,3)-1) + 2 )* dv * 3/8   ! simpson 3/8
     !enddo

     first_call = .false.

  end if


  !-- compute charge density

  tmp(:,:,:)=0.
  
  do ix=1,jx
     do iy=1,jy
        do iz=1,jz
           tmp(iz,iy,ix) = sum(dfu(:,iz,iy,ix,:))*dv
           !do idir=1,jdir
           !   tmp(iz,iy,ix) =  tmp(iz,iy,ix) + sum(dfu(:,iz,iy,ix,idir)*weights)
           !enddo
        end do
     end do
  end do

  call mpi_allreduce(tmp, dens, jz*jy*jx, &
      MPI_REAL8, MPI_SUM, MPI_COMM_WORLD, mpi_err)

      
  !-- find FT of potential --

  dummy3D=dens

  call fft_forward_z(dummy3D)
  call fft_forward_x(dummy3D)
  call fft_forward_y(dummy3D)

  phiFT = dummy3D / kk3D             ! inversing Laplacian in k-space

  do iz=1,jz
     phiFT(iz,:,:) = phiFT(iz,:,:) + phiExtFT(iz)
  end do


  phiFT(1,1,1) = 0.0

  phiFT(:,:,:) = phiFT(:,:,:)*filter(:,:,:)     ! VMD comparibility 
  phiFT(1,:,:) = 0.0                            ! VMD comparibility

 
        
  !-- differentiate potential to find z-component of electric field -- 

  do iz=1,jz
     dummy3D(iz,:,:) = -iunit * kZvec(iz) * phiFT(iz,:,:)
  end do

  ftEz = dummy3D

  call fft_backward_z(dummy3D)
  call fft_backward_x(dummy3D)
  call fft_backward_y(dummy3D)

  Ez3D = dummy3D

  !-- differentiate potential to find x-component of electric field -- 

  do ix=1,jx
     dummy3D(:,:,ix) = -iunit * kXvec(ix) * phiFT(:,:,ix)
  end do

  call fft_backward_z(dummy3D)
  call fft_backward_x(dummy3D)
  call fft_backward_y(dummy3D)

  Ex3D = dummy3D

end subroutine evolve_lin_efield


!*************************************************************************************

!-------------------------------------------------------------------------------------
subroutine rk_rhs(Ux, Ubgk, rhs)
!-------------------------------------------------------------------------------------

  use dimensions
  use parameters, only: iunit, vPhase
  use fft_wrap

  implicit none

  real,    intent(in)    :: Ux(jz, jy, jx, jdir)
  real,    intent(in)    :: Ubgk(jz)
  real,    intent(out)   :: rhs(jz, jy, jx, jdir)
 
  integer                :: iv, ix, iy, iDir
  complex                :: fz(jz)


  iy = jy

  do iDir=1,jDir
     do ix=1, jx
           
        fz(:) = Ux(:, iy, ix, iDir)

        call fft_forward(jz,fz,fz)
        fz = fz/float(jz)
        fz(:) = fz(:) * iunit * kZvec(:) 
        call fft_backward(jz,fz,fz)

        rhs(:, iy, ix, iDir) = - fz(:)* Ubgk(:)
        !rhs(:, iy, ix, iDir) = - fz(:)*( - vPhase )

     enddo
  enddo

end subroutine rk_rhs


!-------------------------------------------------------------------------------------
subroutine evolve_lin_U_long_advect_rk(uxMat, Ubgk, dt)
!-------------------------------------------------------------------------------------

  use dimensions

  implicit none

  real,    intent(inout) :: uxMat(jz, jy, jx, jdir)
  real,    intent(in)    :: Ubgk(jz)
  real,    intent(in)    :: dt


  real                   :: s(jz, jy, jx, jdir)
  real                   :: s0(jz, jy, jx, jdir)
  real                   :: Ux(jz, jy, jx, jdir)
 

  Ux = uxMat;

  call rk_rhs(Ux, Ubgk, s)
  Ux = uxMat + s*dt/2
  s0 = s

  call rk_rhs(Ux, Ubgk, s)
  Ux = uxMat + s*dt/2
  s0 = s0 + 2*s

  call rk_rhs(Ux, Ubgk, s)
  Ux = uxMat + s*dt
  s0 = s0 + 2*s

  call rk_rhs(Ux, Ubgk, s)
  s0 = s0 + s

  uxMat = uxMat + s0*dt/6


end subroutine evolve_lin_U_long_advect_rk


!-------------------------------------------------------------------------------------
subroutine evolve_lin_U_long_advect(uxMat, Ubgk, dt)
!-------------------------------------------------------------------------------------

  use dimensions
  use parameters, only: iunit, vPhase, iWaveFrame
  use fft_wrap

  implicit none

  real,    intent(inout) :: uxMat(jz, jy, jx, jdir)
  real,    intent(in)    :: Ubgk(jz)
  real,    intent(in)    :: dt

  integer                :: iv, ix, iy, iDir
  complex                :: fz(jz)
  complex                :: iVdt


  if ( iWaveFrame .eq. 0 ) return

  !-- advect in z-direction

  iy = jy

  do iDir=1,jDir
     do ix=1, jx
           
        iVdt = iunit * dt * (-vPhase*iWaveFrame)

        fz(:) = uxMat(:, iy, ix, iDir)
        call fft_forward(jz,fz,fz)
        fz = fz/float(jz)
        fz(:) = fz(:) * exp ( -iVdt * kZvec(:) )
        call fft_backward(jz,fz,fz)
        uxMat(:, iy, ix, iDir) = fz(:)

     enddo
  enddo

end subroutine evolve_lin_U_long_advect


!-------------------------------------------------------------------------------------
subroutine evolve_lin_U_trans_advect(uxMat, dt)
!-------------------------------------------------------------------------------------

  use dimensions
  use parameters, only: iunit, ux0
  use fft_wrap

  implicit none

  real,    intent(inout) :: uxMat(jz, jy, jx, jDir)
  real,    intent(in)    :: dt

  integer                :: iv, iz, iy, iDir
  complex                :: fx(jx)
  complex                :: iUdt


  !-- advect in x-direction

  iy=jy

  do iDir=1,jDir
     do iz=1, jz

        iUdt = iunit * dt * ux0(iDir)

        fx(:) = uxMat(iz, iy, :, iDir)
        call fft_forward(jx,fx,fx)
        fx = fx/float(jx)
        fx(:) = fx(:) * exp ( -iUdt * kXvec(:) )
        call fft_backward(jx,fx,fx)
        uxMat(iz, iy, :, iDir) = fx(:)

     enddo
  enddo

end subroutine evolve_lin_U_trans_advect


!-------------------------------------------------------------------------------------
subroutine evolve_lin_U_sourceE( uxMat, Ex2D, dt)
!-------------------------------------------------------------------------------------

  use dimensions

  implicit none

  real,    intent(inout) :: uxMat(jz, jy, jx, jDir)
  real,    intent(in)    :: Ex2D(jz, jy, jx)
  real,    intent(in)    :: dt

  integer                :: iDir

  do iDir=1, jDir
     uxMat(:,:,:, iDir) = uxMat(:,:,:, iDir) + Ex2D(:,:,:)*dt
  end do


end subroutine evolve_lin_U_sourceE


!-------------------------------------------------------------------------------------
subroutine evolve_lin_U_filter(uxMat, filter)
!-------------------------------------------------------------------------------------

  use dimensions
  use fft_wrap
  use fft_3d

  implicit none

  real, intent(inout)   :: uxMat(jz, jy, jx, jDir)
  real,    intent(in)   :: filter(jz,jy,jx)
  complex               :: dummy3D(jz, jy, jx)

  integer               :: iDir


  do iDir=1,jDir

     dummy3D=uxMat(:,:,:,iDir)

     call fft_forward_z(dummy3D)
     call fft_forward_x(dummy3D)
     call fft_forward_y(dummy3D)

     dummy3D(:,:,:) = Dummy3D(:,:,:)*filter(:,:,:)   ! VMD comparibility 
     dummy3D(1,:,:) = 0.0                            ! VMD comparibility

     call fft_backward_z(dummy3D)
     call fft_backward_x(dummy3D)
     call fft_backward_y(dummy3D)

     uxMat(:,:,:,iDir) = dummy3D

  end do

end subroutine evolve_lin_U_filter

!-------------------------------------------------------------------------------------

